import numpy as np

from path import PathButterfly
from double_s import DoubleS


class Trajectory:
    def __init__(self, path: PathButterfly, double_s: DoubleS):
        self.path = path
        self.double_s = double_s

    def get_trajectory(self, t):
        s, ds, dds, ddds = self.double_s.get_trajectory(t)
        p, dp, ddp = self.path.get_path(s)

        x = p
        dx = [dpi * ds for dpi in dp]
        ddx = [dpi * ds + ddpi * dds for dpi, ddpi in zip(dp, ddp)]
        return x, dx, ddx

    def get_trajectory_double_s(self, t):
        return self.double_s.get_trajectory(t)

    def get_t(self):
        return self.double_s.get_total_time()
